"""
【简化版 RAG 流程示例 - 技术详解版 + 两种 BM25 检索对比】

功能：
1️⃣ 文档加载与分块（TextLoader + RecursiveCharacterTextSplitter）
2️⃣ 中文 BM25 关键词检索（两种方式：Langchain封装 vs rank_bm25原生）
3️⃣ FAISS 向量语义检索（HuggingFace Embeddings）
4️⃣ 检索结果融合（RRF算法）
5️⃣ 大模型问答（ChatOpenAI）

技术点说明：
- RAG（Retrieval-Augmented Generation）：结合检索与生成，提高大模型回答精度。
- BM25：基于词频和逆文档频率的检索算法，适合关键词匹配。
- jieba：中文分词工具，解决中文无空格分词问题。
- FAISS：高效向量搜索库，用于语义检索。
- HuggingFace Embeddings：将文本转为向量表示。
- RRF（Reciprocal Rank Fusion）：多路检索结果融合算法，增强召回。
- ChatOpenAI：调用大模型生成最终答案。
"""

from typing import List
import jieba
from dotenv import load_dotenv
from langchain_community.document_loaders import TextLoader
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_ollama import ChatOllama
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter
from rank_bm25 import BM25Okapi

load_dotenv()

# =======================
# 1️⃣ 文档加载与分块
# =======================
loader = TextLoader("RAG_5_Data/data.txt", encoding="utf-8")
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0, separators=['\n'])
docs = text_splitter.split_documents(documents)
texts = [doc.page_content for doc in docs]


# =======================
# 2️⃣ 中文分词函数
# =======================
def jieba_tokenize(text: str) -> list[str]:
    """
    使用 jieba 对中文文本进行分词
    输入: str
    输出: list[str]
    """
    return list(jieba.cut(text))


# =======================
# 3️⃣ BM25 关键词检索 - 两种方式
# =======================

# ----- 方式1：Langchain封装版 BM25Retriever -----
# 优点：
# 1. 封装好，直接集成到 Langchain RAG 流程
# 2. 可直接使用 invoke() 查询
# 缺点：
# 1. 灵活性低，不容易调参数
bm25_lc = BM25Retriever(docs=docs, k=10)
retriever_lc = bm25_lc.from_documents(docs, preprocess_func=jieba_tokenize)
# 测试查询
lc_result = retriever_lc.invoke("如何分布式部署")

# ----- 方式2：rank_bm25 原生 BM25Okapi -----
# 优点：
# 1. 灵活，可完全控制分词、参数
# 2. 中文分词可自由定义
# 缺点：
# 1. 需要手动管理索引和查询
texts_tokens = [jieba_tokenize(t) for t in texts]
bm25_okapi = BM25Okapi(texts_tokens)
okapi_result = bm25_okapi.get_top_n(jieba_tokenize('分布式'), texts, n=10)

# 可以根据需要选择其中一种方式，或者对比结果

# =======================
# 4️⃣ FAISS 向量语义检索
# =======================
embeddings = HuggingFaceEmbeddings(
    model_name="D:/models/BAAIbge-base-zh-v1.5",
    model_kwargs={'device': 'cpu'}
)
faiss_db = FAISS.from_documents(docs, embedding=embeddings)

# =======================
# 5️⃣ RRF 融合函数
# =======================
def rrf(vector_results: List[str], text_results: List[str], k: int = 10, m: int = 60):
    """
    Reciprocal Rank Fusion (RRF) 算法：
    - 作用：融合多路检索结果，提高召回质量
    - 核心原理：按照每个文档在各自检索结果里的排名打分
        score(doc) = Σ(1 / (rank + m))
      其中 rank(doc) 是文档在某一检索结果中的排序（从0开始）
      m 是超参数，防止前几名分数过高
    - 特点：
        1. 排名越靠前，分数越高
        2. 两路检索都有的文档分数会叠加
        3. 不依赖原始相似度分数，只用排名
    - 输入：
        vector_results: 向量检索结果列表
        text_results: BM25检索结果列表
        k: 返回前 k 条融合结果
        m: 超参数
    - 输出：融合排序后的前 k 条文档
    """
    scores = {}
    # 对向量检索结果按排名打分
    for rank, doc in enumerate(vector_results):
        scores[doc] = scores.get(doc, 0) + 1 / (rank + m)
    # 对 BM25 检索结果按排名打分
    for rank, doc in enumerate(text_results):
        scores[doc] = scores.get(doc, 0) + 1 / (rank + m)
    # 返回按分数降序的前 k 条文档
    return [d for d, _ in sorted(scores.items(), key=lambda x: x[1], reverse=True)][:k]


# =======================
# 6️⃣ 查询函数
# =======================
def query(question: str, top_k: int = 5, use_okapi=True):
    """
    RAG 查询函数
    - BM25 + FAISS + RRF 融合
    - use_okapi: True 使用 BM25Okapi，False 使用 BM25Retriever
    """
    # BM25检索
    if use_okapi:
        text_results = bm25_okapi.get_top_n(jieba_tokenize(question), texts, n=top_k)
    else:
        text_results = retriever_lc.invoke(question)

    # FAISS 向量检索
    vector_results_docs = faiss_db.similarity_search(question, k=top_k)
    vector_results = [doc.page_content for doc in vector_results_docs]

    # 融合排序
    fused_results = rrf(vector_results, text_results, k=top_k)
    return fused_results


# =======================
# 7️⃣ 大模型问答
# =======================
#llm = ChatOpenAI(model="xdeepseekv32exp", temperature=0.7)
llm = ChatOllama(temperature=0, model="deepseek-r1:1.5b", base_url="http://127.0.0.1:11434")


def ask_llm(question: str):
    retrieved_docs = query(question, top_k=5, use_okapi=True)
    prompt = f"""
任务目标：根据检索出的文档回答用户问题
任务要求：
    1、不得脱离检索出的文档回答问题
    2、若检索出的文档不包含用户问题的答案，请回答我不知道

用户问题：
{question}

检索出的文档：
{' '.join(retrieved_docs)}
"""
    return llm.invoke(prompt).content


# =======================
# 8️⃣ 测试
# =======================
answer = ask_llm("Java如何处理慢Sql")
print(answer)
